import json
import random
import copy
import os
#用户数据文件夹
user_data_folder = "assets/user_data"
#增强数据保存文件夹
augmented_folder = "assets/user_data/augmented_user_data"
#数据集文件夹
database_file = "assets/database/database.json"
#数据结果文件夹
output_file = "assets/database/main_database.json"

# 定义邻近键映射
nearby_keys = {
    'q': ['w', 'a'],
    'w': ['q', 'e', 'a', 's'],
    'e': ['w', 'r', 's', 'd'],
    'r': ['e', 't', 'd', 'f'],
    't': ['r', 'y', 'f', 'g'],
    'y': ['t', 'u', 'g', 'h'],
    'u': ['y', 'i', 'h', 'j'],
    'i': ['u', 'o', 'j', 'k'],
    'o': ['i', 'p', 'k', 'l'],
    'p': ['o', 'l'],
    'a': ['q', 'w', 's', 'z'],
    's': ['w', 'e', 'a', 'd', 'z', 'x'],
    'd': ['e', 'r', 's', 'f', 'x', 'c'],
    'f': ['r', 't', 'd', 'g', 'c', 'v'],
    'g': ['t', 'y', 'f', 'h', 'v', 'b'],
    'h': ['y', 'u', 'g', 'j', 'b', 'n'],
    'j': ['u', 'i', 'h', 'k', 'n', 'm'],
    'k': ['i', 'o', 'j', 'l', 'm', ','],
    'l': ['o', 'p', 'k', ',', '.'],
    'z': ['a', 's', 'x'],
    'x': ['s', 'd', 'z', 'c'],
    'c': ['d', 'f', 'x', 'v'],
    'v': ['f', 'g', 'c', 'b'],
    'b': ['g', 'h', 'v', 'n'],
    'n': ['h', 'j', 'b', 'm'],
    'm': ['j', 'k', 'n', ','],
    ',': ['k', 'l', 'm', '.'],
    '.': ['l', ','],
    ' ': ['c', 'v', 'b', 'n', 'm', ',', '.']
}

# 数据增强函数
def augment_keystrokes(keystrokes):
    augmented_keystrokes = []
    time_offset = 0.0
    speed_factor = random.uniform(0.8, 1.2)

    for i, keystroke in enumerate(keystrokes):
        if i == 0:
            press_time = 0.0
        else:
            press_time = keystroke['press_time'] * speed_factor + random.uniform(-0.05, 0.05)

        release_time = keystroke['release_time']
        if release_time is not None:
            release_time = release_time * speed_factor + random.uniform(-0.05, 0.05)
        else:
            release_time = press_time + 0.1

        press_time = max(0, press_time)
        release_time = max(press_time + 0.01, release_time)

        press_time = format(press_time, '.4f')
        release_time = format(release_time, '.4f')

        if random.random() < 0.1 and keystroke['key_content'].lower() in nearby_keys:
            key_content = random.choice(nearby_keys[keystroke['key_content'].lower()])
            if keystroke['key_content'].isupper():
                key_content = key_content.upper()
        else:
            key_content = keystroke['key_content']

        augmented_keystrokes.append({
            'key_content': key_content,
            'press_time': float(press_time),
            'release_time': float(release_time)
        })

        if random.random() < 0.2:
            pause_duration = round(random.uniform(0.1, 0.5), 4)
            time_offset += pause_duration

    return augmented_keystrokes

# 数据增强并保存
def augment_and_save_data(input_file, output_folder):
    with open(input_file, 'r') as f:
        original_data = json.load(f)

    augmented_data = copy.deepcopy(original_data)
    original_sequence_count = len(original_data)
    new_sequence_count = original_sequence_count * 9

    for _ in range(new_sequence_count):
        original_entry = random.choice(original_data)
        new_entry = {
            'subject': original_entry['subject'],
            'keystrokes': augment_keystrokes(original_entry['keystrokes'])
        }
        augmented_data.append(new_entry)

    base_filename = os.path.basename(input_file)
    output_filename = os.path.join(output_folder, base_filename)
    with open(output_filename, 'w') as f:
        json.dump(augmented_data, f, indent=4)

    print(f"数据增强完成，已保存到 {output_filename}")

# 合并数据
def merge_data(augmented_folder, database_file, output_file):
    merged_data = []

    # 加载原始 database.json 数据
    if os.path.exists(database_file):
        with open(database_file, 'r') as f:
            merged_data = json.load(f)

    # 遍历 augmented_user_data 文件夹中的所有 JSON 文件
    for filename in os.listdir(augmented_folder):
        if filename.endswith('.json'):
            filepath = os.path.join(augmented_folder, filename)
            with open(filepath, 'r') as f:
                data = json.load(f)
                merged_data.extend(data)

    # 保存合并后的数据到 main_database.json
    with open(output_file, 'w') as f:
        json.dump(merged_data, f, indent=4)

    print(f"数据合并完成，已保存到 {output_file}")

# 主函数
def main():

    # 确保 augmented_user_data 文件夹存在
    if not os.path.exists(augmented_folder):
        os.makedirs(augmented_folder)

    # 遍历 user_data 文件夹中的所有 JSON 文件并进行数据增强
    for filename in os.listdir(user_data_folder):
        if filename.endswith('.json'):
            input_file = os.path.join(user_data_folder, filename)
            augment_and_save_data(input_file, augmented_folder)

    # 合并增强后的数据到 main_database.json
    merge_data(augmented_folder, database_file, output_file)

if __name__ == "__main__":
    main()